/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Description: Define api of stream cache consumer implementation.
 */

#ifndef DATASYSTEM_CLIENT_STREAM_CACHE_CONSUMER_IMPL_H
#define DATASYSTEM_CLIENT_STREAM_CACHE_CONSUMER_IMPL_H

#include "datasystem/client/stream_cache/client_base_impl.h"
#include "datasystem/client/stream_cache/producer_consumer_worker_api.h"
#include "datasystem/common/eventloop/timer_queue.h"
#include "datasystem/common/shared_memory/shm_unit.h"
#include "datasystem/common/stream_cache/cursor.h"
#include "datasystem/common/util/queue/circular_queue.h"
#include "datasystem/stream/element.h"
#include "datasystem/utils/optional.h"
#include "datasystem/utils/status.h"

namespace datasystem {
namespace client {
namespace stream_cache {
class __attribute((visibility("default"))) ConsumerImpl : public ClientBaseImpl {
public:
    /**
     * @brief Construct Consumer.
     * @param[in] streamName The name of the stream.
     * @param[in] config The configure of subscription, such as subscription name, subscription mode.
     * @param[in] consumerId The uuid of the consumer which is generated by Worker.
     * @param[in] rsp The response from rpc request
     * @param[in] workerApi Used to call worker service through rpc.
     * @param[in] mmapManager Used to receive and mmap fd that passed from worker.
     * @param[in] listenWorker Listening to the worker survival status.
     * @param[in] autoAck Toggles if auto ack is enabled or not for this consumer
     */
    ConsumerImpl(std::string streamName, std::string tenantId, SubscriptionConfig config, std::string consumerId,
                 const SubscribeRspPb &rsp, std::shared_ptr<client::stream_cache::ProducerConsumerWorkerApi> workerApi,
                 std::shared_ptr<StreamClientImpl> client, client::MmapManager *mmapManager,
                 std::shared_ptr<client::ListenWorker> listenWorker, bool autoAck);

    ~ConsumerImpl() override;

    /**
     * @brief Initialize the ConsumerImpl
     * @return status of the call.
     */
    Status Init() override;

    /**
     * @brief Sets the state of the Producer to CLOSE
     */
    void SetInactive() override;

    /**
     * @brief Get expectNum elements form the subscription.
     * @param[in] expectNum Parameter to indicate the number of expected elements to receive. If this optional
     * parameter is not set, it will receive any amount.
     * @param[in] timeoutMs The timeout millisecond of elements to be Receive.
     * @param[out] outElements The received elements to be read.
     * @return K_OK on success; the error code otherwise.
     *         K_UNKNOWN_ERROR: it's up to return message.
     *         K_NOT_FOUND: the id of stream is not found.
     *         K_INVALID: invalid parameter.
     *         K_RPC_UNAVAILABLE: didn't receive any response from server.
     *         K_DUPLICATED: the consumer already had pending receive.
     *         K_SC_PRODUCER_NOT_FOUND: one or more producer in the stream are dead
     *         K_SC_STREAM_IN_RESET_STATE: stream currently in reset state
     *         K_SC_ALREADY_CLOSED: the consumer is already closed
     */
    Status Receive(Optional<uint32_t> expectNum, uint32_t timeoutMs, std::vector<Element> &outElements);

    /**
     * @brief Acknowledge elements that had been read by this consumer.
     * @param[in] elementId The element id that to be acknowledged.
     * @return K_OK on success; the error code otherwise.
     *         K_UNKNOWN_ERROR: it's up to return message.
     *         K_NOT_FOUND: the id of stream is not found.
     *         K_INVALID: invalid parameter.
     *         K_SC_STREAM_IN_RESET_STATE: stream currently in reset state.
     *         K_SC_ALREADY_CLOSED: the consumer is already closed
     */
    Status Ack(uint64_t elementId);

    /**
     * @brief Close the consumer, after close it will not allow Receive and Ack Elements.
     * Calling Close() on an already closed consumer will return K_OK.
     * @return K_OK on success; the error code otherwise.
     *         K_UNKNOWN_ERROR: it's up to return message.
     *         K_NOT_FOUND: the id of stream is not found.
     *         K_INVALID: invalid parameter.
     *         K_RUNTIME_ERROR: delete sub node in global scope fail on master process.
     *         K_SC_STREAM_IN_RESET_STATE: stream currently in reset state.
     */
    Status Close();

    /**
     * @brief Reset the consumer by cleaning up data and metadata.
     * @return K_OK on success; the error code otherwise.
     *         K_RUNTIME_ERROR: Consumer is not in reset State.
     */
    Status Reset();

    /**
     * @brief Set consumer state to Reset. Cancel prefetch of data. During reset it will not allow Receive, Ack.
     * @return K_OK on success; the error code otherwise.
     *         K_RUNTIME_ERROR: consumer is in close state already.
     */
    Status SetStateToReset();

    /**
     * @brief Resume the consumer, allowing Receive, Ack and Close again.
     * @return K_OK on success; the error code otherwise.
     *         K_RUNTIME_ERROR: consumer is not in reset state.
     */
    Status Resume();

    /**
     * @brief Get prefix of log entry.
     * @param[in] withSubName Whether print subscription name, by default false.
     */
    std::string LogPrefix() const override;

    /**
     * @brief Get the amount of received elements since this consumer construct, and the amount of elements
     * not processed.
     * @param[out] totalElements the amount of elements received by this consumer.
     * @param[out] notProcessedElements the amount of elements not processed (received but not ack-ed).
     */
    void GetStatisticsMessage(uint64_t &totalElements, uint64_t &notProcessedElements);

    /**
     * @brief Get the consumer id
     * @return Consumer id.
     */
    const std::string &GetConsumerId() const
    {
        return consumerId_;
    }

private:
    /**
     * @brief Check the state_ of consumer and return status.
     * @return Status of the call.
     */
    Status CheckNormalState() const override;

    /**
     * @brief Get expectNum elements form the local cache queue.
     * @param[in] expectNum The number of elements to be read.
     * @param[out] outElements The received elements to be read.
     * @return Status of the call.
     */
    Status GetElementFromCache(uint32_t expectNum, std::vector<Element> &outElements);

    /**
     * @brief When the cache is sufficient to satisfy the receive, return these cached elements.
     * @param[in] cacheLength The length of the cache
     * @param[in] receiveNum The number of elements to fetch
     * @param[out] outElements The fetched elements output
     */
    Status CacheFetch(uint32_t cacheLength, uint32_t receiveNum, std::vector<Element> &outElements);

    /**
     * @brief Check if the cache can satisfy a receive and drive the cache fetch if needed. Also, invoke cache prefetch
     * (if configured for it).
     * @param[in/out] cacheLength The length of the cache.
     * @param[in] receiveNum The number of elements to fetch
     * @param[in] timeoutMs The timeout for receives
     * @param[out] outElements The fetched elements output
     * @param[out] needsRecv T/F if the cache handler informs that a receive will be needed.
     */
    Status CacheHandler(uint32_t &cacheLength, uint32_t receiveNum, uint32_t timeoutMs,
                        std::vector<Element> &outElements, bool &needsRecv);

    /**
     * @brief Called from the receive codepath, this function decides if it should send an Ack or not and also gathers
     * some statistics about Ack logistics to help decide on future Acks
     * @param[in] fetchedElements The element list that was just recently fetched from the worker
     * @return status of the call
     */
    Status PostRecvAckHandler(const std::vector<Element> &fetchedElements);

    /**
     * @brief Called from the receive codepath, this function sends Ack before receive if AutoAck is enabled
     */
    void PreRecvAckHandler();

    /**
     * @brief Get the last ack cursor from the work area
     * @return last ack cursor
     */
    uint64_t GetWALastAckCursor() const
    {
        return cursor_->GetWALastAckCursor();
    }

    /**
     * @brief Update the last ack cursor in work area
     * @param elementId
     */
    void UpdateWALastAckCursor(uint64_t elementId) const
    {
        cursor_->UpdateWALastAckCursor(elementId);
    }

    /**
     * @brief Internal function to replenish the local client cache by calling receive to the worker.
     * @return status of the call
     */
    Status PrefetchReceive();

    /**
     * @brief Get expectNum elements form the subscription.
     * @param[in] expectNum Parameter to indicate the number of expected elements to receive. If this optional
     * parameter is not set, it will receive any amount.
     * @param[in] timeoutMs The timeout millisecond of elements to be Receive.
     * @param[out] outElements The received elements to be read.
     * @return K_OK on success; the error code otherwise.
     *         K_UNKNOWN_ERROR: it's up to return message.
     *         K_NOT_FOUND: the id of stream is not found.
     *         K_INVALID: invalid parameter.
     *         K_RPC_UNAVAILABLE: didn't receive any response from server.
     *         K_DUPLICATED: the consumer already had pending receive.
     *         K_SC_PRODUCER_NOT_FOUND: one or more producer in the stream are dead
     *         K_SC_STREAM_IN_RESET_STATE: stream currently in reset state
     */
    Status ReceiveImpl(Optional<uint32_t> expectNum, uint32_t timeoutMs, std::vector<Element> &outElements);

    Status PrefetchEntry(uint32_t targetNum, uint32_t timeoutMs);
    Status GetDataPage(const ShmView &shmView, std::shared_ptr<StreamDataPage> &out);
    Status PrefetchElements(uint32_t timeoutMs, std::shared_ptr<StreamDataPage> &dataPage, uint32_t targetNum,
                            uint32_t &totalFetched, bool nonBlockingFetch);
    Status GetPrefetchPage(int64_t timeoutMs, std::shared_ptr<StreamDataPage> &out);
    Status LocatePrefetchPage(int64_t timeoutMs, std::shared_ptr<StreamDataPage> &out);
    Status ReleasePage(std::shared_ptr<StreamDataPage> &page) const;
    Status ExtractBigElements(std::shared_ptr<StreamDataPage> &dataPage, std::vector<DataElement> &recvElements);

    /**
     * @brief Logs the received elements, consumed elements, and acked (processed) elements
     */
    void LogConsumerCursors();

    /**
     * @brief Extract header's version from element.
     * @param[in] element Element (header + data).
     * @param[out] version The header's version.
     * @return K_OK on success.
     */
    Status ExtractVersion(DataElement &element, ElementHeader::Version &version);

    /**
     * @brief Extract headers from elements and process each header according to the header version.
     * @param[in] elements Element (header + data)
     * @return K_OK on success.
     */
    Status ProcessHeaders(std::vector<DataElement> &elements);

    /**
     * @brief Verify element is received in order for its producer.
     * @param[in] recvElement Element's data
     * @param[in] recvHeader Element's header
     * @return K_OK on success
     */
    Status VerifyElement(const DataElement &recvElement, const ElementHeader &recvHeader);

    // for make_shared to access private/protected constructor.
    friend std::shared_ptr<Consumer> std::make_shared<Consumer>();

    const SubscriptionConfig config_;
    const std::string consumerId_;
    std::atomic<uint64_t> lastRecvCursor_;
    std::atomic<uint64_t> pageBoundaryCursor_;
    std::atomic<int64_t> bytesSinceLastAck_{ 0 };
    std::atomic<uint64_t> consumedElements_{ 0 };
    float avgEleSize_{ 0 };
    uint64_t lastAvgCount_{ 0 };
    uint32_t cachePrefetchLWM_{ 0 };
    std::atomic_uint64_t ackedElementId_;  // The ID of last Ack'ed element.
    mutable std::shared_timed_mutex idxMutex_;
    std::map<uint64_t, std::shared_ptr<StreamDataPage>> idx_;
    std::shared_ptr<StreamDataPage> lastPage_;  // last page. ref count > 0
    const SubscribeRspPb rsp_;

    std::unique_ptr<HeapCircularQ<Element>> elementCacheQueue_;
    bool autoAck_{ false };
    std::future<Status> prefetchStatus_;
    WaitPost receiveWp_;

    // recording the last element's seqNo from a particicular producer, any normal update should be +1 only.
    std::unordered_map<std::string, DataVerificationHeader::SeqNo> producerLastSeqNoReceive_;
    FirstCallTracer recvTracer_;
};
}  // namespace stream_cache
}  // namespace client
}  // namespace datasystem

#endif  // DATASYSTEM_CLIENT_STREAM_CACHE_CONSUMER_IMPL_H
